import os
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain.output_parsers import BooleanOutputParser,ResponseSchema,StructuredOutputParser,JsonOutputToolsParser
from langchain_core.output_parsers import StrOutputParser,JsonOutputParser
# 优先使用langchain_core下而非langchain的，因为其是核心标准实现，而不带core是更高层的封装，同一子包路径下两者不可混用，当确实差功能时，再选择统一非core版本
from langchain_core.prompts import PromptTemplate,ChatPromptTemplate,HumanMessagePromptTemplate,SystemMessagePromptTemplate
from langchain_core.runnables import RunnableLambda


# 最简单chain
def simple_chain():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 使用管道符将llm的输出给到输出解析器，该解析器为字符串解析器，所以最终输出为字符串
    chain = model | StrOutputParser()
    question = "用中文介绍一下你自己"
    result = chain.invoke(question)
    print(f"!!! simple chain result: {result}")

# 带有提示词的chain，更加稳定控制llm输出
def simple_chain_with_prompt():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    system_prompt="""
    你是一个乐意助人的助手，请根据用户的问题给出回答
    """
    human_prompt="""
    这是用户的问题： {user_query}， 请用 yes 或 no 来回答
    """
    system_prompt_template = SystemMessagePromptTemplate.from_template(system_prompt)
    human_prompt_template = HumanMessagePromptTemplate.from_template(human_prompt)
    prompt_template = ChatPromptTemplate.from_messages([system_prompt_template, human_prompt_template])
    chain = prompt_template | model | StrOutputParser()
    result = chain.invoke({"user_query": "1+1=2吗？"})
    print(f"!!! simple chain with prompt result: {result}")

# 使用Bool结构化解析器固定输出为 True/False
def simple_chain_with_bool_parser():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    system_prompt = """
        你是一个乐意助人的助手，请根据用户的问题给出回答
        """
    human_prompt = """
        这是用户的问题： {user_query}， 请用 yes 或 no 来回答
        """
    system_prompt_template = SystemMessagePromptTemplate.from_template(system_prompt)
    human_prompt_template = HumanMessagePromptTemplate.from_template(human_prompt)
    prompt_template = ChatPromptTemplate.from_messages([system_prompt_template, human_prompt_template])
    chain = prompt_template | model | BooleanOutputParser()
    result = chain.invoke({"user_query": "1+1=2吗？"})
    print(f"!!! simple chain with boolean parser result: {result}")

# 使用结构化解析器固定输出为{}字典类型
def address_extract_chain():
    schemas = [
        ResponseSchema(name="country", description="地址中的国家"),
        ResponseSchema(name="city", description="地址中的城市"),
        ResponseSchema(name="area", description="地址中的区域"),
        ResponseSchema(name="street", description="地址中的街道"),
        ResponseSchema(name="zipcode", description="地址中的邮编"),
    ]
    output_parser = StructuredOutputParser.from_response_schemas(schemas)
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    system_prompt = """
            你是一个地址解析助手，请根据用户的输入，提取出地址中的国家、城市、区域、街道、邮编等信息，并按要求的格式输出{format_instructions}。
            """
    human_prompt = """
            这是用户输入的地址： {input_address}
            """
    system_prompt_template = SystemMessagePromptTemplate.from_template(system_prompt)
    human_prompt_template = HumanMessagePromptTemplate.from_template(human_prompt)
    prompt_template = ChatPromptTemplate.from_messages([system_prompt_template, human_prompt_template])
    chain = prompt_template.partial(format_instructions=output_parser.get_format_instructions()) | model | output_parser
    result = chain.invoke({"input_address": "中国四川省成都市金牛区人民北路二段1号 610031"})
    print(f"!!! address extract chain result: {result}, type of result: {type(result)}")


# 使用JSON解析器固定输出为json
def user_info_extract_chain():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 在prompt中定义json需要对{}进行转义，变成{{}}，否则会被当成占位符
    system_prompt = """
                你是一个用户信息提取助手，需要提取用户的名字，年龄，性别，并且以json格式输出[{{"name":"A","age":10,"gender":"男"}}]。
                """
    human_prompt = """
                这是用户输入的用户信息： {input_user_info}
                """
    system_prompt_template = SystemMessagePromptTemplate.from_template(system_prompt)
    human_prompt_template = HumanMessagePromptTemplate.from_template(human_prompt)
    prompt_template = ChatPromptTemplate.from_messages([system_prompt_template, human_prompt_template])
    chain = prompt_template | model | JsonOutputParser()
    result = chain.invoke({"input_user_info": "张三19岁是个女生"})
    print(f"!!! user info extract chain result: {result}, type of result: {type(result)}")

# 创建复合链，多个chain之间可以组合
def generate_news_then_summary_combined_chain():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    news_generate_prompt_msg = ("请根据以下新闻标题撰写一段简短的新闻内容（100字以内）："
              "标题：{title}")
    # 构建简单的prompt，此prompt不指定角色，仅用于简单的输入输出场景
    news_generate_prompt = PromptTemplate.from_template(news_generate_prompt_msg)
    news_generate_chain = news_generate_prompt | model

    schemas = [
        ResponseSchema(name="time", description="事件发生的时间"),
        ResponseSchema(name="location", description="事件发生的地点"),
        ResponseSchema(name="event", description="发生的具体事件"),
    ]
    parser = StructuredOutputParser.from_response_schemas(schemas)
    news_summary_prompt_msg = ("请从下面这段新闻内容中提取关键信息："
                               "{news}"
                               "格式为：{format_instructions}")
    news_summary_prompt = PromptTemplate.from_template(news_summary_prompt_msg)

    news_summary_chain = news_summary_prompt.partial(format_instructions=parser.get_format_instructions()) | model | parser

    # 组装chain
    full_chain = news_generate_chain | news_summary_chain
    result = full_chain.invoke({"title": "苹果公司在加州发布新款AI芯片"})
    print(f"!!! news generate and summary chain result: {result}, type of result: {type(result)}")


# Langchain支持自定义节点，可用于将chain的中间结果打印出来，方便调试
def generate_news_then_summary_combined_chain_with_debug():
    def debug_print(text):
        print(f"生成的新闻内容: {text}")
        return text

    debug_node = RunnableLambda(debug_print)

    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    news_generate_prompt_msg = ("请根据以下新闻标题撰写一段简短的新闻内容（100字以内）："
              "标题：{title}")
    # 构建简单的prompt，此prompt不指定角色，仅用于简单的输入输出场景
    news_generate_prompt = PromptTemplate.from_template(news_generate_prompt_msg)
    news_generate_chain = news_generate_prompt | model

    schemas = [
        ResponseSchema(name="time", description="事件发生的时间"),
        ResponseSchema(name="location", description="事件发生的地点"),
        ResponseSchema(name="event", description="发生的具体事件"),
    ]
    parser = StructuredOutputParser.from_response_schemas(schemas)
    news_summary_prompt_msg = ("请从下面这段新闻内容中提取关键信息："
                               "{news}"
                               "格式为：{format_instructions}")
    news_summary_prompt = PromptTemplate.from_template(news_summary_prompt_msg)

    news_summary_chain = news_summary_prompt.partial(format_instructions=parser.get_format_instructions()) | model | parser

    # 组装chain
    full_chain = news_generate_chain | debug_node | news_summary_chain
    result = full_chain.invoke({"title": "苹果公司在加州发布新款AI芯片"})
    print(f"!!! news generate and summary chain result: {result}, type of result: {type(result)}")



if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    print(f"DEEPSEEK_API_KEY: {DEEPSEEK_API_KEY}")
    simple_chain()
    simple_chain_with_prompt()
    simple_chain_with_bool_parser()
    address_extract_chain()
    user_info_extract_chain()
    generate_news_then_summary_combined_chain()
    generate_news_then_summary_combined_chain_with_debug()







